// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/array/array_dict.h>

#include <algorithm>
#include <climits>
#include <cstdint>
#include <limits>
#include <utility>
#include <vector>

#include <nebula/array/array_primitive.h>
#include <nebula/array/data.h>
#include <nebula/array/dict_internal.h>
#include <nebula/array/util.h>
#include <nebula/core/buffer.h>
#include <nebula/core/chunked_array.h>
#include <nebula/compute/api.h>
#include <nebula/core/datum.h>
#include <turbo/utility/status.h>
#include <nebula/core/table.h>
#include <nebula/types/type.h>
#include <nebula/types/type_traits.h>
#include <nebula/bits/bit_util.h>
#include <nebula/bits/bitmap_ops.h>
#include <turbo/base/checked_cast.h>
#include <nebula/numeric/int_util.h>
#include <turbo/log/logging.h>
#include <nebula/core/visit_type_inline.h>

namespace nebula {

    using internal::CopyBitmap;

    // ----------------------------------------------------------------------
    // DictionaryArray

    const std::shared_ptr<Array> &DictionaryArray::indices() const { return indices_; }

    int64_t DictionaryArray::get_value_index(int64_t i) const {
        const uint8_t *indices_data = data_->buffers[1]->data();
        // If the value is non-negative then we can use the unsigned path
        switch (indices_->type_id()) {
            case Type::UINT8:
            case Type::INT8:
                return static_cast<int64_t>(indices_data[data_->offset + i]);
            case Type::UINT16:
            case Type::INT16:
                return static_cast<int64_t>(
                        reinterpret_cast<const uint16_t *>(indices_data)[data_->offset + i]);
            case Type::UINT32:
            case Type::INT32:
                return static_cast<int64_t>(
                        reinterpret_cast<const uint32_t *>(indices_data)[data_->offset + i]);
            case Type::UINT64:
            case Type::INT64:
                return static_cast<int64_t>(
                        reinterpret_cast<const uint64_t *>(indices_data)[data_->offset + i]);
            default:
                KCHECK(false) << "unreachable";
                return -1;
        }
    }

    DictionaryArray::DictionaryArray(const std::shared_ptr<ArrayData> &data)
            : dict_type_(turbo::checked_cast<const DictionaryType *>(data->type.get())) {
        KCHECK_EQ(data->type->id(), Type::DICTIONARY);
        KCHECK_NE(data->dictionary, nullptr);
        set_data(data);
    }

    void DictionaryArray::set_data(const std::shared_ptr<ArrayData> &data) {
        this->Array::set_data(data);
        auto indices_data = data_->copy();
        indices_data->type = dict_type_->index_type();
        indices_data->dictionary = nullptr;
        indices_ = make_array(indices_data);
    }

    DictionaryArray::DictionaryArray(const std::shared_ptr<DataType> &type,
                                     const std::shared_ptr<Array> &indices,
                                     const std::shared_ptr<Array> &dictionary)
            : dict_type_(turbo::checked_cast<const DictionaryType *>(type.get())) {
        KCHECK_EQ(type->id(), Type::DICTIONARY);
        KCHECK_EQ(indices->type_id(), dict_type_->index_type()->id());
        KCHECK_EQ(dict_type_->get_value_type()->id(), dictionary->type()->id());
                DKCHECK(dict_type_->get_value_type()->equals(*dictionary->type()));
        auto data = indices->data()->copy();
        data->type = type;
        data->dictionary = dictionary->data();
        set_data(data);
    }

    const std::shared_ptr<Array> &DictionaryArray::dictionary() const {
        if (!dictionary_) {
            // TODO(GH-36503) this isn't thread safe
            dictionary_ = make_array(data_->dictionary);
        }
        return dictionary_;
    }

    turbo::Result<std::shared_ptr<Array>> DictionaryArray::from_arrays(
            const std::shared_ptr<DataType> &type, const std::shared_ptr<Array> &indices,
            const std::shared_ptr<Array> &dictionary) {
        if (type->id() != Type::DICTIONARY) {
            return turbo::failed_precondition_error("Expected a dictionary type");
        }
        const auto &dict = turbo::checked_cast<const DictionaryType &>(*type);
        if (indices->type_id() != dict.index_type()->id()) {
            return turbo::failed_precondition_error(
                    "Dictionary type's index type does not match "
                    "indices array's type");
        }
        TURBO_RETURN_NOT_OK(internal::CheckIndexBounds(*indices->data(),
                                                       static_cast<uint64_t>(dictionary->length())));
        return std::make_shared<DictionaryArray>(type, indices, dictionary);
    }

    bool DictionaryArray::can_compare_indices(const DictionaryArray &other) const {
                DKCHECK(dictionary()->type()->equals(other.dictionary()->type()))
                << "dictionaries have differing type " << *dictionary()->type() << " vs "
                << *other.dictionary()->type();

        if (!indices()->type()->equals(other.indices()->type())) {
            return false;
        }

        auto min_length = std::min(dictionary()->length(), other.dictionary()->length());
        return dictionary()->range_equals(other.dictionary(), 0, min_length, 0);
    }

    // ----------------------------------------------------------------------
    // Dictionary transposition

    namespace {

        inline bool IsTrivialTransposition(const int32_t *transpose_map,
                                           int64_t input_dict_size) {
            for (int64_t i = 0; i < input_dict_size; ++i) {
                if (transpose_map[i] != i) {
                    return false;
                }
            }
            return true;
        }

        turbo::Result<std::shared_ptr<ArrayData>> TransposeDictIndices(
                const std::shared_ptr<ArrayData> &data, const std::shared_ptr<DataType> &in_type,
                const std::shared_ptr<DataType> &out_type,
                const std::shared_ptr<ArrayData> &dictionary, const int32_t *transpose_map,
                MemoryPool *pool) {
            // Note that in_type may be different from data->type if data is of type ExtensionType
            if (in_type->id() != Type::DICTIONARY || out_type->id() != Type::DICTIONARY) {
                return turbo::failed_precondition_error("Expected dictionary type");
            }
            const int64_t in_dict_len = data->dictionary->length;
            const auto &in_dict_type = turbo::checked_cast<const DictionaryType &>(*in_type);
            const auto &out_dict_type = turbo::checked_cast<const DictionaryType &>(*out_type);

            const auto &in_index_type = *in_dict_type.index_type();
            const auto &out_index_type =
                    turbo::checked_cast<const FixedWidthType &>(*out_dict_type.index_type());

            if (in_index_type.id() == out_index_type.id() &&
                IsTrivialTransposition(transpose_map, in_dict_len)) {
                // Index type and values will be identical => we can simply reuse
                // the existing buffers.
                auto out_data =
                        ArrayData::create(out_type, data->length, {data->buffers[0], data->buffers[1]},
                                          data->null_count, data->offset);
                out_data->dictionary = dictionary;
                return out_data;
            }

            // Default path: compute a buffer of transposed indices.
            TURBO_MOVE_OR_RAISE(
                    auto out_buffer,
                    allocate_buffer(data->length * (out_index_type.bit_width() / CHAR_BIT), pool));

            // Shift null buffer if the original offset is non-zero
            std::shared_ptr<Buffer> null_bitmap;
            if (data->offset != 0 && data->null_count != 0) {
                TURBO_MOVE_OR_RAISE(null_bitmap, CopyBitmap(pool, data->buffers[0]->data(),
                                                            data->offset, data->length));
            } else {
                null_bitmap = data->buffers[0];
            }

            auto out_data = ArrayData::create(out_type, data->length,
                                              {null_bitmap, std::move(out_buffer)}, data->null_count);
            out_data->dictionary = dictionary;
            TURBO_RETURN_NOT_OK(internal::TransposeInts(
                    in_index_type, out_index_type, data->get_values<uint8_t>(1, 0),
                    out_data->GetMutableValues<uint8_t>(1, 0), data->offset, out_data->offset,
                    data->length, transpose_map));
            return out_data;
        }

        struct CompactTransposeMapVisitor {
            const std::shared_ptr<ArrayData> &data;
            nebula::MemoryPool *pool;
            std::unique_ptr<Buffer> output_map;
            std::shared_ptr<Array> out_compact_dictionary;

            template<typename IndexArrowType>
            turbo::Status CompactTransposeMapImpl() {
                int64_t index_length = data->length;
                int64_t dict_length = data->dictionary->length;
                if (dict_length == 0) {
                    output_map = nullptr;
                    out_compact_dictionary = nullptr;
                    return turbo::OkStatus();
                } else if (index_length == 0) {
                    TURBO_MOVE_OR_RAISE(out_compact_dictionary,
                                        MakeEmptyArray(data->dictionary->type, pool));
                    TURBO_MOVE_OR_RAISE(output_map, allocate_buffer(0, pool));
                    return turbo::OkStatus();
                }

                using CType = typename IndexArrowType::c_type;
                const CType *indices_data = data->get_values<CType>(1);
                std::vector<bool> dict_used(dict_length, false);
                CType dict_len = static_cast<CType>(dict_length);
                int64_t dict_used_count = 0;
                for (int64_t i = 0; i < index_length; i++) {
                    if (data->is_null(i)) {
                        continue;
                    }

                    CType current_index = indices_data[i];
                    if (current_index < 0 || current_index >= dict_len) {
                        return turbo::out_of_range_error(
                                "Index out of bounds while compacting dictionary array: ", current_index,
                                "(dictionary is ", dict_length, " long) at position ", i);
                    }
                    if (dict_used[current_index]) continue;
                    dict_used[current_index] = true;
                    dict_used_count++;

                    if (dict_used_count == dict_length) {
                        // The dictionary is already compact, so just return here
                        output_map = nullptr;
                        out_compact_dictionary = nullptr;
                        return turbo::OkStatus();
                    }
                }

                using BuilderType = NumericBuilder<IndexArrowType>;
                using nebula::compute::Take;
                using nebula::compute::TakeOptions;
                BuilderType dict_indices_builder(pool);
                TURBO_RETURN_NOT_OK(dict_indices_builder.Reserve(dict_used_count));
                TURBO_MOVE_OR_RAISE(output_map,
                                    allocate_buffer(dict_length * sizeof(int32_t), pool));
                auto *output_map_raw = output_map->mutable_data_as<int32_t>();
                int32_t current_index = 0;
                for (CType i = 0; i < dict_len; i++) {
                    if (dict_used[i]) {
                        dict_indices_builder.unsafe_append(i);
                        output_map_raw[i] = current_index;
                        current_index++;
                    } else {
                        output_map_raw[i] = -1;
                    }
                }
                TURBO_MOVE_OR_RAISE(std::shared_ptr<nebula::Array> compacted_dict_indices,
                                    dict_indices_builder.finish());
                TURBO_MOVE_OR_RAISE(auto compacted_dict_res,
                                    Take(Datum(data->dictionary), compacted_dict_indices,
                                         TakeOptions::NoBoundsCheck()));
                out_compact_dictionary = compacted_dict_res.make_array();
                return turbo::OkStatus();
            }

            template<typename Type>
            enable_if_integer<Type, turbo::Status> Visit(const Type &) {
                return CompactTransposeMapImpl<Type>();
            }

            turbo::Status Visit(const DataType &type) {
                return turbo::failed_precondition_error("Expected an Index Type of Int or UInt");
            }
        };

        turbo::Result<std::unique_ptr<Buffer>> CompactTransposeMap(
                const std::shared_ptr<ArrayData> &data, MemoryPool *pool,
                std::shared_ptr<Array> &out_compact_dictionary) {
            if (data->type->id() != Type::DICTIONARY) {
                return turbo::failed_precondition_error("Expected dictionary type");
            }

            const auto &dict_type = turbo::checked_cast<const DictionaryType &>(*data->type);
            CompactTransposeMapVisitor visitor{data, pool, nullptr, nullptr};
            TURBO_RETURN_NOT_OK(visit_type_inline(*dict_type.index_type(), &visitor));

            out_compact_dictionary = visitor.out_compact_dictionary;
            return std::move(visitor.output_map);
        }
    }  // namespace

    turbo::Result<std::shared_ptr<Array>> DictionaryArray::transpose(
            const std::shared_ptr<DataType> &type, const std::shared_ptr<Array> &dictionary,
            const int32_t *transpose_map, MemoryPool *pool) const {
        TURBO_MOVE_OR_RAISE(auto transposed,
                            TransposeDictIndices(data_, data_->type, type, dictionary->data(),
                                                 transpose_map, pool));
        return make_array(std::move(transposed));
    }

    turbo::Result<std::shared_ptr<Array>> DictionaryArray::compact(MemoryPool *pool) const {
        std::shared_ptr<Array> compact_dictionary;
        TURBO_MOVE_OR_RAISE(std::unique_ptr<Buffer> transpose_map,
                            CompactTransposeMap(this->data_, pool, compact_dictionary));

        if (transpose_map == nullptr) {
            return std::make_shared<DictionaryArray>(this->data_);
        } else {
            return this->transpose(this->type(), compact_dictionary,
                                   transpose_map->data_as<int32_t>(), pool);
        }
    }

    // ----------------------------------------------------------------------
    // Dictionary unification

    namespace {

        template<typename T>
        class DictionaryUnifierImpl : public DictionaryUnifier {
        public:
            using ArrayType = typename TypeTraits<T>::ArrayType;
            using DictTraits = typename internal::DictionaryTraits<T>;
            using MemoTableType = typename DictTraits::MemoTableType;

            DictionaryUnifierImpl(MemoryPool *pool, std::shared_ptr<DataType> value_type)
                    : pool_(pool), value_type_(value_type), memo_table_(pool) {}

            turbo::Status unify(const Array &dictionary, std::shared_ptr<Buffer> *out) override {
                if (dictionary.null_count() > 0) {
                    return turbo::invalid_argument_error("Cannot yet unify dictionaries with nulls");
                }
                if (!dictionary.type()->equals(*value_type_)) {
                    return turbo::invalid_argument_error("Dictionary type different from unifier: ",
                                                         dictionary.type()->to_string());
                }
                const ArrayType &values = turbo::checked_cast<const ArrayType &>(dictionary);
                if (out != nullptr) {
                    TURBO_MOVE_OR_RAISE(auto result,
                                        allocate_buffer(dictionary.length() * sizeof(int32_t), pool_));
                    auto result_raw = reinterpret_cast<int32_t *>(result->mutable_data());
                    for (int64_t i = 0; i < values.length(); ++i) {
                        TURBO_RETURN_NOT_OK(memo_table_.GetOrInsert(values.get_view(i), &result_raw[i]));
                    }
                    *out = std::move(result);
                } else {
                    for (int64_t i = 0; i < values.length(); ++i) {
                        int32_t unused_memo_index;
                        TURBO_RETURN_NOT_OK(memo_table_.GetOrInsert(values.get_view(i), &unused_memo_index));
                    }
                }
                return turbo::OkStatus();
            }

            turbo::Status unify(const Array &dictionary) override { return unify(dictionary, nullptr); }

            turbo::Status get_result(std::shared_ptr<DataType> *out_type,
                                     std::shared_ptr<Array> *out_dict) override {
                int64_t dict_length = memo_table_.size();
                std::shared_ptr<DataType> index_type;
                if (dict_length <= std::numeric_limits<int8_t>::max()) {
                    index_type = int8();
                } else if (dict_length <= std::numeric_limits<int16_t>::max()) {
                    index_type = int16();
                } else if (dict_length <= std::numeric_limits<int32_t>::max()) {
                    index_type = int32();
                } else {
                    index_type = int64();
                }
                // Build unified dictionary type with the right index type
                *out_type = nebula::dictionary(index_type, value_type_);

                // Build unified dictionary array
                TURBO_MOVE_OR_RAISE(
                        auto data, DictTraits::GetDictionaryArrayData(pool_, value_type_, memo_table_,
                                                                      0 /* start_offset */));
                *out_dict = make_array(data);
                return turbo::OkStatus();
            }

            turbo::Status get_result_with_index_type(const std::shared_ptr<DataType> &index_type,
                                                 std::shared_ptr<Array> *out_dict) override {
                Int64Scalar dict_length(memo_table_.size());
                if (!internal::IntegersCanFit(dict_length, *index_type).ok()) {
                    return turbo::invalid_argument_error(
                            "These dictionaries cannot be combined.  The unified dictionary requires a "
                            "larger index type.");
                }

                // Build unified dictionary array
                TURBO_MOVE_OR_RAISE(
                        auto data, DictTraits::GetDictionaryArrayData(pool_, value_type_, memo_table_,
                                                                      0 /* start_offset */));
                *out_dict = make_array(data);
                return turbo::OkStatus();
            }

        private:
            MemoryPool *pool_;
            std::shared_ptr<DataType> value_type_;
            MemoTableType memo_table_;
        };

        struct MakeUnifier {
            MemoryPool *pool;
            std::shared_ptr<DataType> value_type;
            std::unique_ptr<DictionaryUnifier> result;

            MakeUnifier(MemoryPool *pool, std::shared_ptr<DataType> value_type)
                    : pool(pool), value_type(value_type) {}

            template<typename T>
            enable_if_no_memoize<T, turbo::Status> Visit(const T &) {
                // Default implementation for non-dictionary-supported datatypes
                return turbo::unimplemented_error("Unification of ", *value_type,
                                                  " dictionaries is not implemented");
            }

            template<typename T>
            enable_if_memoize<T, turbo::Status> Visit(const T &) {
                result.reset(new DictionaryUnifierImpl<T>(pool, value_type));
                return turbo::OkStatus();
            }
        };

        struct RecursiveUnifier {
            MemoryPool *pool;

            // Return true if any of the arrays was changed (including descendents)
            turbo::Result<bool> unify(std::shared_ptr<DataType> type, ArrayDataVector *chunks) {
                        DKCHECK(!chunks->empty());
                bool changed = false;
                std::shared_ptr<DataType> ext_type = nullptr;

                if (type->id() == Type::EXTENSION) {
                    ext_type = std::move(type);
                    type = turbo::checked_cast<const ExtensionType &>(*ext_type).storage_type();
                }

                // unify all child dictionaries (if any)
                if (type->num_fields() > 0) {
                    ArrayDataVector children(chunks->size());
                    for (int i = 0; i < type->num_fields(); ++i) {
                        std::transform(chunks->begin(), chunks->end(), children.begin(),
                                       [i](const std::shared_ptr<ArrayData> &array) {
                                           return array->child_data[i];
                                       });
                        TURBO_MOVE_OR_RAISE(bool child_changed,
                                            unify(type->field(i)->type(), &children));
                        if (child_changed) {
                            // Only do this when unification actually occurred
                            for (size_t j = 0; j < chunks->size(); ++j) {
                                (*chunks)[j]->child_data[i] = std::move(children[j]);
                            }
                            changed = true;
                        }
                    }
                }

                // unify this dictionary
                if (type->id() == Type::DICTIONARY) {
                    const auto &dict_type = turbo::checked_cast<const DictionaryType &>(*type);
                    // XXX Ideally, we should unify dictionaries nested in value_type first,
                    // but DictionaryUnifier doesn't supported nested dictionaries anyway,
                    // so this will fail.
                    TURBO_MOVE_OR_RAISE(auto unifier,
                                        DictionaryUnifier::create(dict_type.get_value_type(), this->pool));
                    // unify all dictionary array chunks
                    BufferVector transpose_maps(chunks->size());
                    for (size_t j = 0; j < chunks->size(); ++j) {
                                DKCHECK_NE((*chunks)[j]->dictionary, nullptr);
                        TURBO_RETURN_NOT_OK(
                                unifier->unify(*make_array((*chunks)[j]->dictionary), &transpose_maps[j]));
                    }
                    std::shared_ptr<Array> dictionary;
                    TURBO_RETURN_NOT_OK(unifier->get_result_with_index_type(dict_type.index_type(), &dictionary));
                    for (size_t j = 0; j < chunks->size(); ++j) {
                        TURBO_MOVE_OR_RAISE(
                                (*chunks)[j],
                                TransposeDictIndices(
                                        (*chunks)[j], type, type, dictionary->data(),
                                        reinterpret_cast<const int32_t *>(transpose_maps[j]->data()), this->pool));
                        if (ext_type) {
                            (*chunks)[j]->type = ext_type;
                        }
                    }
                    changed = true;
                }

                return changed;
            }
        };

    }  // namespace

    turbo::Result<std::unique_ptr<DictionaryUnifier>> DictionaryUnifier::create(
            std::shared_ptr<DataType> value_type, MemoryPool *pool) {
        MakeUnifier maker(pool, value_type);
        TURBO_RETURN_NOT_OK(visit_type_inline(*value_type, &maker));
        return std::move(maker.result);
    }

    turbo::Result<std::shared_ptr<ChunkedArray>> DictionaryUnifier::unify_chunked_array(
            const std::shared_ptr<ChunkedArray> &array, MemoryPool *pool) {
        if (array->num_chunks() <= 1) {
            return array;
        }

        ArrayDataVector data_chunks(array->num_chunks());
        std::transform(array->chunks().begin(), array->chunks().end(), data_chunks.begin(),
                       [](const std::shared_ptr<Array> &array) { return array->data(); });
        TURBO_MOVE_OR_RAISE(bool changed,
                            RecursiveUnifier{pool}.unify(array->type(), &data_chunks));
        if (!changed) {
            return array;
        }
        ArrayVector chunks(array->num_chunks());
        std::transform(data_chunks.begin(), data_chunks.end(), chunks.begin(),
                       [](const std::shared_ptr<ArrayData> &data) { return make_array(data); });
        return std::make_shared<ChunkedArray>(std::move(chunks), array->type());
    }

    turbo::Result<std::shared_ptr<Table>> DictionaryUnifier::unify_table(const Table &table,
                                                                        MemoryPool *pool) {
        ChunkedArrayVector columns = table.columns();
        for (auto &col: columns) {
            TURBO_MOVE_OR_RAISE(col, DictionaryUnifier::unify_chunked_array(col, pool));
        }
        return Table::create(table.schema(), std::move(columns), table.num_rows());
    }

}  // namespace nebula
